#!/bin/bash
#set -ex
set -e
# -------------------------------------------------------------------------------
# Filename:     benchmark4vllm_70b_300_tp8.sh
# UpdateDate:   2025/02/08
# Description:  一键运行性能测试，并保存性能测试日志及测试期间的cnmon信息到各自log。
# Example:      ./benchmark4vllm_70b_300_tp8.sh
# Depends:
# Notes:        本脚本适用于docker容器环境下运行；
# -------------------------------------------------------------------------------
#切换到推理环境并设置环境变量
#source /torch/venv3/pytorch_infer/bin/activate
export MODEL_PATH=/data/models/llm/models/DeepSeek-R1-Distill-Llama-70B
export MODEL_Name=DeepSeek-R1-Distill-Llama-70B
export MLU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export MLU_TP_NUM=8
CMD_TIME=$(date +%Y%m%d%H%M%S.%N)
# 进入工作目录
CUR_DIR=$(cd $(dirname $0);pwd)
pushd "${CUR_DIR}/.."
if [ ! -d "./log" ]; then mkdir -p "./log";fi
popd
export VLLM_LATENCY_DEBUG=true
# 由于驱动的原因，需要export CN_TASKTOPO_RESIDENT=FALSE
#export CN_TASKTOPO_RESIDENT=FALSE
export USE_PAGED=false
#设置模型路径
#MODEL_Name=DeepSeek-R1-Distill-Llama-70B
MODEL_FULLNAME=${MODEL_PATH}
#-block-size 设置 BLOCKSIZE;
# paged 模式仅支持默认值16, 在非paged模式默认值为2048;
#非paged 模式 --block-size 必须大于等于 input_seqlen + output_seqlen
export MAX_TOTAL_TOKENS=512
block_size=${MAX_TOTAL_TOKENS}
#max-num-batched-tokens 必须是 MAX_TOTAL_TOKENS 或 block-size 的整数倍。
#设置最大长度
max_model_len=8192
array_bs=(1 2 4)
array_seqlen_input=(512 1024 2048 4096 8192 16384 32768)
array_seqlen_output=(256)

for seqlen_input in "${array_seqlen_input[@]}"; do
    for seqlen_output in "${array_seqlen_output[@]}"; do
        #不同条件不同赋值测试,减少部分测试
        case $seqlen_input in
            512)
                array_bs=(1 8 16 32 64 128)
                ;;
            1024)
                array_bs=(1 4 8 16 32 64 128)
                ;;
            2048)
                array_bs=(1 8 16 32 64)
                ;;
            4096)
                array_bs=(1 4 8 16 32)
                ;;
            8192)
                array_bs=(1 8 16)
                ;;
            16384)
                array_bs=(1 4 8)
                ;;
            32768)
                array_bs=(1 2 4)
                ;;
            *)
                array_bs=(1 2 4)
                ;;
        esac
        for bs in "${array_bs[@]}"; do
            for tp_nums in ${MLU_TP_NUM}; do
                # 0. 设置参数
                export MAX_TOTAL_TOKENS=$(expr ${seqlen_input} + ${seqlen_output})
                block_size=${MAX_TOTAL_TOKENS}
                max_num_batched_tokens=$(expr ${seqlen_input} \* $bs)
                if [ $max_model_len -gt $max_num_batched_tokens ]; then
                    max_num_batched_tokens=$max_model_len
                fi
                LOG_FILENAME="${CUR_DIR}/../log/${MODEL_Name}_latency_blks${block_size}_sli${seqlen_input}_slo${seqlen_output}_tp${tp_nums}_bs${bs}_${CMD_TIME}"
                # 1. 监控 cnmon 命令并将输出追加到日志文件中
                while true; do cnmon -c ${MLU_VISIBLE_DEVICES} >> "${LOG_FILENAME}_cnmon.log";sleep 1;done 2>&1 &
                PID_CNMONProcess=$!
                echo "PID_CNMONProcess: $PID_CNMONProcess"
                # 2. 性能测试
                cd /workspace/vllm-*
                python benchmarks/benchmark_latency.py \
                        --model ${MODEL_FULLNAME} \
                        --tokenizer ${MODEL_FULLNAME} \
                        --num-iters-warmup 1 \
                        --num-iters 3 \
                        --dtype float16 \
                        --input-len ${seqlen_input} \
                        --output-len ${seqlen_output} \
                        --block-size ${block_size} \
                        -tp ${tp_nums} \
                        --trust-remote-code \
                        --max-model-len ${max_model_len} \
                        --max-num-batched-tokens ${max_num_batched_tokens} \
                        --gpu-memory-utilization 0.98 \
                        --enforce-eager \
                        --batch-size ${bs} 2>&1 | tee "${LOG_FILENAME}.log"
                # 3. 删除【实时记录cnmon信息】的进程;# 使用 tail -f ./*_cnmon.log 可以实时查看日志文件的内容
                sleep 3
                kill -9 $PID_CNMONProcess
            done
        done
    done
done
sleep 1
#切换目录
pushd "${CUR_DIR}/.." && ls -la log/